#include <stdio.h>
#include <stdlib.h>
#include <time.h>

/* want 64-bit Words */
#define Word unsigned long long
/* want 32-bit Indexes */
#define Index unsigned long
/* want 8 bit Counts */
#define Count unsigned char
#define COUNTBINS (1 << (sizeof(Count) * 8))
typedef enum { DESCRIP, FASTA } State;
/* typedef enum { A=0, C=1, G=2, T=3 } Base; */
#define A (0ULL)
#define C (1ULL)
#define G (2ULL)
#define T (3ULL)
#ifndef BASESUNUSED
#define BASESUNUSED (0) 
#endif
#define BITSPERBASE (2UL)
#define BITSUNUSED (BASESUNUSED*BITSPERBASE)
#define BITSPERWORD (sizeof(Word) * 8UL)
#define BASEMASK ((1UL << BITSPERBASE) - 1)
#define BASESPERBYTE (8/BITSPERBASE)
#define BASESPERWORD (BITSPERWORD/BITSPERBASE - BASESUNUSED)

/* #define HASHSIZE (357913931) */
#ifndef HASHSIZE
#define HASHSIZE (101000777)
#endif
Word *hash; /* [HASHSIZE]; */
Count *count; /* [HASHSIZE] = {0}; */
Index reads = 0; /* global read count */
Index nmers = 0; /* global nmer count */
Index distinct = 0;
clock_t begin;

/*
 *  Don't call bases twice in same function call because it uses a static char array.
 */
#define shiftin(w, r, b, g, p) { \
      w <<= (BITSPERBASE + BITSUNUSED);  \
      w >>= BITSUNUSED; \
      w |= b;\
      r >>= BITSPERBASE; r |= (~b & BASEMASK) << (BITSPERWORD - (BITSPERBASE + BITSUNUSED)); \
      g++; \
      if (g >= (BASESPERWORD-BASESUNUSED)) { \
        Word tw; \
        tw = (w <= r)? w : r; \
        addcount(tw, p); \
      } \
}

/*
 *  Forward declarations
 */
Word revcomp(Word w);
char *bases(Word w);
void addcount(Word w, int proc);

int main(int argc, char **argv) {

	printf("Hi there!");

  char pc, c;
  State s = FASTA;
  Word nmer = 0;
  Word rev = 0;
  int ngood = 0;
  Index i;
  Index countbins[COUNTBINS] = {0};
  int proc;

  /* printf("arg: %s\n", argv[1]); */
#ifndef UNIBATCH
  sscanf(argv[1], "%d", &proc);
#endif

  begin = clock();
  printf("Hash size is: %ld\n", HASHSIZE);
  fflush(stdout);

  hash = (Word *) malloc(sizeof(Word) * HASHSIZE);
  count = (Count *) calloc(sizeof(Count), HASHSIZE);

  /* printf("count vals: %d %d %d %d\n", count[100], count[3000], count[30000], count[2<<16 -1]); */

  while ((c = getchar()) != EOF) {
    /* fprintf(stderr, "Testing: %c\n", c); */
    if (s == DESCRIP) { 
      if (c == '\n') {
	s = FASTA;
	ngood = 0;
      }
      /* else ignore */
    }
    else /* s == FASTA */ {
      switch (c) {
        case '>': {
	  s = DESCRIP;
	  reads++;
	  if (! (reads & 0xFFFF)) {
	    fprintf(stderr, "\r%ld reads, %ld nmers, %ld distinct in hash (%ld percent), %ld seconds.", 
		    reads, nmers, distinct, ((distinct * 100UL)/nmers),
		    (long) ((clock() - begin)/CLOCKS_PER_SEC));
	    fflush(stderr);
	  }
	  break;
	}
        case 'a': case 'A': {
	  shiftin(nmer, rev, A, ngood, proc);
	  break;
	}
        case 'c': case 'C': {
	  shiftin(nmer, rev, C, ngood, proc);
	  break;
	}
        case 'g': case 'G': {
	  shiftin(nmer, rev, G, ngood, proc);
	  break;
	}
        case 't': case 'T': {
	  shiftin(nmer, rev, T, ngood, proc);
	  break;
	}
        case ' ': case '\t': case '\n': case '*': {
	  /* ignore */
	  break;
	}
        default: {
	  nmer = 0;
	  rev = 0;
	  ngood = 0;
	}
      }
    }
  }
  fprintf(stderr, "\n");
  printf("Total reads:\t%ld\tnmers:\t%ld\tdistinct:\t%ld\tseconds:\t%ld\n",
	 reads, nmers, distinct, (long) ((clock() - begin)/CLOCKS_PER_SEC));
  for (i = 0; i < HASHSIZE; i++) {
    countbins[count[i]]++;
#ifndef NOMERS
    if (count[i] > 8) {
      printf("%s\t%4d\n", bases(hash[i]), count[i]);
    }
#endif
  }
  printf("Nmer abundance bins:\n");
  for (i = 0; i < COUNTBINS; i++) {
    printf("%d\t%d\n", i, countbins[i]);
  }
}

Word revcomp(Word w) {
  int i;
  Word rev = 0;
  /* 
   * The word w is treated as a string of 2-bit bases
   * (A == 01, T == 11, C == 01, G == 10)
   * So long as we reverse the string basewise,
   * we just need to bit-complement to get the final answer.
   */
  for (i = 1; i <= BASESPERWORD; i++) {
    /*
     *  Shift result left and add in last base of input; then advance input
     */
    rev = (rev << BITSPERBASE) + (w & BASEMASK);
    w >>= BITSPERBASE;
  }
  return ~rev; /* bitwise complement == basewise complement */
}

char *bases(Word w) {
  static char b[BASESPERWORD+1];
  char base = '\0';
  int basebits;
  int i;
  for (i = (BASESPERWORD - BASESUNUSED); i >= 0; i--) {
    b[i] = base;
    basebits = w & BASEMASK;
    w >>= BITSPERBASE;
    switch (basebits) {
      case A: 
	base = 'a';
	break;
      case C: 
	base = 'c';
	break;
      case G: 
	base = 'g';
	break;
      case T: 
	base = 't';
	break;
    }
  }
  return b;
}

void addcount(Word w, int proc) {
  Index start, probe, step, i;
  Count t;
  Word temp;
  Count bitcount = 0;

  Index up, lo;
  lo = w & 0xFFFFFFFF;
  up = w >> 32;

#ifndef UNIBATCH
  for (temp = lo ^ up; temp; temp >>= 5) {
    bitcount ^= temp & 0x1F;
  }
  if (bitcount != proc) {
    return;
  }
#endif
  nmers++;

  probe = start = (lo ^ up) % HASHSIZE;
  step = (w & 0xF) + 1; /* 1..16 */
  while (count[probe] && hash[probe] != w) {
    probe += step;
    if (probe >= HASHSIZE) probe -= HASHSIZE;
    if (probe == start) break;
  }
  if (! count[probe]) {
    distinct++;
    count[probe] = 1;
    hash[probe] = w;
  }
  else if (hash[probe] != w) {
    fprintf(stderr, "\nFull hash on word %s, probe %ld, insertion %ld\n", 
	    bases(w), probe, distinct);
    exit(1);
  }
  else if ((t = count[probe] + 1) > 0) {
    count[probe] = t;
  }
}
